package mongo

import (
	"container/heap"
	"core/utils"
	"sync"
	"time"

	"github.com/globalsign/mgo"
)

type (
	dialContext struct {
		sync.Mutex
		sessions sessionHeap
	}
	MongodbConf struct {
		AppName      string
		MongodbConn  string
		SessionCount int
	}
)

var (
	appname    string
	mongodbCtx *dialContext
)

func InitMongodb(conf *MongodbConf) (closeCtx func(), err error) {
	db, err := dialMongodb(conf.MongodbConn, conf.AppName)
	if err != nil {
		utils.Error(err)
		return
	}
	mongodbCtx = initMongodbContext(db, conf.SessionCount, mgo.Strong)
	closeCtx = func() {
		mongodbCtx.close()
	}
	return
}

func dialMongodb(uri, dbName string) (mongoDB *mgo.Database, err error) {
	return dialMongodbWithTimeout(uri, dbName, 10*time.Second, 30*time.Second, 5*time.Minute)
}

func dialMongodbWithTimeout(url, dbName string, dialTimeout, SyncTimeout, timeout time.Duration) (mongoDB *mgo.Database, err error) {
	s, err := mgo.DialWithTimeout(url, dialTimeout) // 创建连接超时时间
	if err != nil {
		return nil, err
	}
	s.SetSyncTimeout(SyncTimeout) // 同步查询超时时间
	s.SetSocketTimeout(timeout)   // socket超时时间
	mongoDB = s.DB(dbName)
	return
}

func initMongodbContext(mongoDB *mgo.Database, sessionNum int, consistency mgo.Mode) (context *dialContext) {
	context = new(dialContext)
	context.sessions = make(sessionHeap, sessionNum)
	for i := 0; i < sessionNum; i++ {
		session := mongoDB.Session.New()
		session.SetMode(consistency, true)
		context.sessions[i] = &mongoSession{session, 0, i}
	}
	heap.Init(&context.sessions)
	return
}

func Execute(collectionName string, block func(c *mgo.Collection) (err error)) (err error) {
	c, closeFunc := mongodbCtx.c(collectionName)
	defer closeFunc()
	return block(c)
}

// ref get session from pool
func (ctx *dialContext) ref() *mongoSession {
	ctx.Lock()
	s := ctx.sessions[0]
	if s.ref == 0 {
		s.Refresh()
	}
	s.ref++
	heap.Fix(&ctx.sessions, 0)
	ctx.Unlock()
	return s
}

// unRef return session to pool
func (ctx *dialContext) unRef(s *mongoSession) {
	ctx.Lock()
	s.ref--
	heap.Fix(&ctx.sessions, s.index)
	ctx.Unlock()
}

// Close close all session
func (ctx *dialContext) close() {
	ctx.Lock()
	for _, s := range ctx.sessions {
		s.Close()
		if s.ref != 0 {
			utils.Error("session ref = ", s.ref)
		}
	}
	ctx.Unlock()
}

// MARK:- collection
// c get mgo.Collection from pool
func (ctx *dialContext) c(collection string) (c *mgo.Collection, closedb func()) {
	defer utils.CatchPanic()
	s := ctx.ref()
	c = s.DB(appname).C(collection)
	closedb = func() { ctx.unRef(s) }
	return
}
